from pathlib import Path
import numpy as np
import pyrealsense2 as rs
from ezrobot.sensor.base import BaseSensor
# from line_profiler import LineProfiler
from ezrobot.utils.general import accurate_sleep

class RealSense(BaseSensor):
    def get_device_serial_numbers(self):
        devices = rs.context().devices
        return [d.get_info(rs.camera_info.serial_number) for d in devices]

    def __init__(self, index="021422060263", dim=(640, 480), fps=30, depth=True, exposure=40000): #相机好像只能到30
        assert index in self.get_device_serial_numbers()
        self.index = index
        self.depth = depth
        self.pipe = rs.pipeline()
        self.cfg = rs.config()
        self.cfg.enable_device(self.index)
        self.cfg.enable_stream(rs.stream.color, dim[0], dim[1], rs.format.bgr8, fps)
        if self.depth:
            self.cfg.enable_stream(rs.stream.depth, dim[0], dim[1], rs.format.z16, fps)
        self.profile = self.pipe.start(self.cfg)
        self.s = self.profile.get_device().first_depth_sensor()
        self.depth_scale = self.s.get_depth_scale()
        self.s.set_option(rs.option.exposure, exposure)

        # Create an align object
        # rs.align allows us to perform alignment of depth frames to others frames
        # The "align_to" is the stream type to which we plan to align depth frames.
        align_to = rs.stream.color
        self.align = rs.align(align_to)

    def recv_data(self):
        frames = self.pipe.wait_for_frames()
        aligned_frames = self.align.process(frames)
        color_frame = aligned_frames.get_color_frame()
        if self.depth:
            depth_frame = aligned_frames.get_depth_frame()

        if color_frame.is_video_frame():
            image = np.asarray(color_frame.get_data())
            if self.depth and depth_frame.is_depth_frame():
                depth = np.expand_dims(np.asarray(depth_frame.get_data()), axis=2)
                depth_meters = depth * self.depth_scale  # shape不变，单位为米，dtype=float32/float64
                return {"color": image, "depth": depth_meters}
            else:
                return {"color": image}
        else:
            accurate_sleep(0.05)
            return self.recv_data()

    def close(self):
        self.pipe.stop()
        self.cfg.disable_all_streams()


if __name__ == "__main__":
    import cv2
    realsense = RealSense()
    while 1:
        data_dict = realsense.recv_data()
        cv2.imshow(f"color img", data_dict['color'])
        accurate_sleep(0.1)
        if cv2.waitKey(1) == 27:  # ESC退出
            break
    realsense.close()